'''
4.编排器-工作期（Orchestrator-Worker） 定义：一个"编排器"负责任务分解与调度，多个"工作期"（Worker）执行具体子任务。

应用场景：
AI代理系统：Orchestrator 拆解用户请求 → 分配给"搜索Worker""写作Worker""代码Worker"等。
自动化工作流平台：如低代码AI平台，Orchestrator管理流程，Worker执行具体API调用或模型推理。
企业级AI服务：统一入口接收请求，Orchestrator协调多个微服务（如NLP、CV、语音）协同工作。
复杂任务自动化：如"生成一份市场分析报告" → 分解为数据收集、分析、图表生成、报告撰写等任务。
✅ 优势：模块化、可扩展、支持异构系统集成。
'''

# 导入必要的库和模块
from typing import Annotated, List
import operator
from pydantic import BaseModel, Field
from langchain_core.messages import HumanMessage, SystemMessage
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.types import Send
from utils import init_llm, set_env, visualize_graph


class Section(BaseModel):
    """
    报告章节的数据模型
    用于定义报告的各个部分的结构
    """
    name: str = Field(
        description="报表中此部分的名称。",
    )
    description: str = Field(
        description="本节将涵盖的主要主题和概念的简要概述。",
    )


class Sections(BaseModel):
    """
    报告章节集合的数据模型
    包含所有报告章节的列表
    """
    sections: List[Section] = Field(
        description="报告的各个部分。",
    )


class State(TypedDict):
    """
    工作流状态的数据模型
    定义了整个报告生成过程中需要维护的状态信息
    """
    topic: str  # 报告主题
    sections: list[Section]  # 报告章节列表
    completed_sections: Annotated[
        list, operator.add
    ]  # 已完成的章节内容，使用operator.add进行合并
    final_report: str  # 最终生成的完整报告


class WorkerState(TypedDict):
    """
    工作节点状态的数据模型
    定义了单个工作节点处理时的状态信息
    """
    section: Section  # 当前要处理的章节
    completed_sections: Annotated[list, operator.add]  # 已完成的章节内容


def orchestrator(state: State):
    """
    编排器函数 - 负责任务分解和规划
    
    功能：
    1. 接收报告主题
    2. 使用LLM生成报告章节计划
    3. 返回章节列表供后续工作节点处理
    
    Args:
        state: 当前工作流状态
        
    Returns:
        包含章节列表的字典
    """

    # 使用LLM生成报告章节计划
    report_sections = planner.invoke(
        [
            SystemMessage(content="生成报表的计划。"),
            HumanMessage(content=f"以下是报告主题：{state['topic']}"),
        ]
    )

    return {"sections": report_sections.sections}


def llm_call(state: WorkerState):
    """
    工作节点函数 - 负责编写单个报告章节
    
    功能：
    1. 接收章节信息（名称和描述）
    2. 使用LLM生成该章节的详细内容
    3. 返回完成的章节内容
    
    Args:
        state: 工作节点状态，包含要处理的章节信息
        
    Returns:
        包含完成章节内容的字典
    """

    # 使用LLM生成章节内容
    section = llm.invoke(
        [
            SystemMessage(
                content="按照提供的名称和描述编写报告部分。每个部分不包括序言。使用 Markdown 格式。"
            ),
            HumanMessage(
                content=f"这是部分名称： {state['section'].name} 还有描述: {state['section'].description}"
            ),
        ]
    )

    return {"completed_sections": [section.content]}


def synthesizer(state: State):
    """
    综合器函数 - 将所有完成的章节合并成最终报告
    
    功能：
    1. 收集所有已完成的章节内容
    2. 使用分隔符将它们合并
    3. 生成最终的完整报告
    
    Args:
        state: 当前工作流状态
        
    Returns:
        包含最终报告的字典
    """

    # 获取所有已完成的章节
    completed_sections = state["completed_sections"]

    # 使用分隔符合并所有章节内容
    completed_report_sections = "\n\n---\n\n".join(completed_sections)

    return {"final_report": completed_report_sections}


def assign_workers(state: State):
    """
    工作分配函数 - 为每个章节分配工作节点
    
    功能：
    1. 遍历所有章节
    2. 为每个章节创建一个工作节点任务
    3. 返回任务列表供并行执行
    
    Args:
        state: 当前工作流状态
        
    Returns:
        工作节点任务列表
    """

    # 为每个章节创建一个工作节点任务
    return [Send("llm_call", {"section": s}) for s in state["sections"]]


def build_graph():
    """
    构建工作流图函数
    
    功能：
    1. 创建状态图实例
    2. 添加所有节点（编排器、工作节点、综合器）
    3. 定义节点之间的连接关系
    4. 编译并返回可执行的工作流图
    
    Returns:
        编译后的工作流图
    """
    
    # 创建状态图实例
    orchestrator_worker_builder = StateGraph(State)

    # 添加各个节点
    orchestrator_worker_builder.add_node("orchestrator", orchestrator)  # 编排器节点
    orchestrator_worker_builder.add_node("llm_call", llm_call)  # 工作节点
    orchestrator_worker_builder.add_node("synthesizer", synthesizer)  # 综合器节点

    # 定义工作流程：开始 -> 编排器 -> 工作节点 -> 综合器 -> 结束
    orchestrator_worker_builder.add_edge(START, "orchestrator")  # 从开始到编排器
    orchestrator_worker_builder.add_conditional_edges(
        "orchestrator", assign_workers, ["llm_call"]  # 从编排器到工作节点（条件边）
    )
    orchestrator_worker_builder.add_edge("llm_call", "synthesizer")  # 从工作节点到综合器
    orchestrator_worker_builder.add_edge("synthesizer", END)  # 从综合器到结束

    # 编译工作流图
    graph = orchestrator_worker_builder.compile()

    return graph


if __name__ == "__main__":
    # 主程序入口
    set_env()  # 设置环境变量
    llm = init_llm(temperature=0)  # 初始化LLM，温度设为0以获得确定性输出
    planner = llm.with_structured_output(Sections)  # 创建结构化输出的规划器
    graph = build_graph()  # 构建工作流图
 
    # visualize_graph(graph)  # 可视化工作流图（已注释）

    # 执行工作流，生成关于LLM缩放定律的报告
    state = graph.invoke({"topic": "创建有关 LLM 缩放定律的报告"})
    print(state["final_report"])  # 打印最终生成的报告







